#version 450

layout(std430, binding = 0) readonly buffer Octree { uint data[]; } octree;

#include "octree.glsl"

layout(binding = 1) uniform CameraUbo {
    vec4 position;
	vec4 look;
	vec4 right;
	vec4 up;
} cameraUbo;

layout(binding = 2) uniform sampler2D beamImage;

layout(location = 0) out vec4 outColor;

layout(push_constant) uniform PushConstants {
	uint width;
	uint height;
	uint beamEnable;
	uint beamSize;
	uint viewType;
} pushConstants;

vec3 CameraGenRay(vec2 coord) {
	coord = coord * 2.0f - 1.0f;
	return normalize(cameraUbo.look.xyz - cameraUbo.right.xyz * coord.x - cameraUbo.up.xyz * coord.y);
}

vec3 Heat(in float x) { return sin(clamp(x, 0.0, 1.0) * 3.0 - vec3(1, 2, 3)) * 0.5 + 0.5; }

void main() {
	vec3 o = cameraUbo.position.xyz;
	vec3 d = CameraGenRay(ivec2(gl_FragCoord.xy) / vec2(pushConstants.width, pushConstants.height));

	float beam;
	if (pushConstants.beamEnable == 1) {
		ivec2 beamCoord = ivec2(gl_FragCoord.xy / pushConstants.beamSize);
		beam = min(min(texelFetch(beamImage, beamCoord, 0).r, texelFetch(beamImage, beamCoord + ivec2(1, 0), 0).r),
		           min(texelFetch(beamImage, beamCoord + ivec2(0, 1), 0).r,
		               texelFetch(beamImage, beamCoord + ivec2(1, 1), 0).r));
		o += d * beam;
	}

	vec3 pos, color, normal;
	uint iter;
	bool hit = OctreeRayMarchLeaf(o, d, pos, color, normal, iter);
	if (!hit) {
		pos = vec3(1.0);
		normal = vec3(0.0);
		color = vec3(1.0);
	}

	switch(pushConstants.viewType)
	{
		case 0:
			outColor = vec4(pow(color, vec3(1.0 / 2.2)), 1.0);
		break;
		case 1:
			outColor = vec4(normal * 0.5 + 0.5, 1.0);
		break;
		case 2:
			outColor = vec4(pos - 1.0, 1.0);
		break;
		case 3:
			outColor = vec4(Heat(float(iter) * 1.0 / 128.0), 1.0);
		break;
		case 4:
			outColor = vec4(vec3(texelFetch(beamImage, ivec2(gl_FragCoord.xy / pushConstants.beamSize), 0).r), 1.0);
		break;
		case 5:
			outColor = vec4(d, 1.0);
		break;
	}
}
